import { DataNode } from "antd/es/tree"
import { toPinyin } from "./common"

export type Key = string | number
export interface TreeItem<K extends Key = Key> {
    key: K
    value: K
    title?: string
    label?: string
    parent_key?: K
    children?: TreeItem<K>[]
}
export interface TreeItemMapKey<O> extends Partial<Record<keyof TreeItem, keyof O | string>> {
    __with_pinyin__?: boolean
    /** 保留源对象中其余属性 */
    __with_rest__?: boolean
}
export type MapTreeDataResult<K extends Key> = readonly [TreeItem<K>[], Map<K, TreeItem<K>>]

export const mapTreeData = function <T extends Key, O>(options: O[] = [], {
    key = 'key',
    value = 'value',
    title = 'title',
    label = 'label',
    parent_key = 'parent_key',
    children = 'children',
    __with_pinyin__ = false,
    __with_rest__ = false,
}: TreeItemMapKey<O>) {
    let with_children = false
    const children_map = new Map<T, (TreeItem<T> & O)[]>()
    const item_map = new Map<T, (TreeItem<T> & O)>()
    function loop (options: O[], pid?: T) {
        const result: Array<TreeItem<T> & O> = []
        for (let i = 0; i < options.length; i++) {
            const item = options[i];
            const {
                [key]: _key,
                [value]: _value,
                [title]: _title,
                [label]: _label,
                [parent_key]: _parent_key,
                [children]: _children,
                ...rest
            } = item as Record<string, string | number | O[]>
            const obj = {
                ...(__with_rest__ ? rest : {}),
                key: _key,
                value: _value || _key,
                title: _title,
                label: _label || _title,
                parent_key: _parent_key,
                children: _children ? loop(_children as O[], _key as T) : undefined,
            } as unknown as TreeItem<T> & O
            if (__with_pinyin__) {
                obj.label = `${obj.label} ${toPinyin(obj.label)}`
            }
            if (typeof pid != 'undefined') {
                // 如果有pid说明存在children，就不在进行整理
                with_children = true
                obj.parent_key = pid
            } else if (typeof obj.parent_key != 'undefined') {
                // 如果有parent_key 需要整理一下children数组
                const _children_ = children_map.get(obj.parent_key) || []
                children_map.set(obj.parent_key, _children_.concat(obj))
            } else {
                // 顶层不处理
            }
            item_map.set(obj.key, obj)
            result.push(obj)
        }
        return result.length > 0 ? result : undefined
    }
    let tree = loop(options)
    if (!with_children) {
        tree = tree?.filter(o => {
            o.children = children_map.get(o.key)
            // 严格检查，避免服务端返回部分数据导致用parentid检测根节点错误
            return typeof o.parent_key === 'undefined' || !item_map.has(o.parent_key)
        })
    }
    return [tree || [], item_map] as const
}

export const searchTreeData = function (
    keyword = '',
    treeData: MapTreeDataResult<Key>,
    filter?: (item: TreeItem) => boolean,
    render = (item: TreeItem): DataNode => item,
): TreeItem[] {
    keyword = keyword.trim()
    const visibles = new Set()
    const map = treeData[1]
    map.forEach((v, k) => {
        if (visibles.has(k)) {
            return
        }
        
        if (
            (!keyword || v.label.includes(keyword)) &&
            (!filter || filter(v))
        ) {
            do {
                visibles.add(v.key)
                v = map.get(v.parent_key)
            } while (v)
        }
    })
    function loop (items: TreeItem[]) {
        return items?.map(item => {
            if (visibles.has(item.key)) {
                return {
                    ...render(item),
                    children: loop(item.children)
                }
            }
        }).filter(a => a)
    }
    return loop(treeData[0])
}